from contextvars import ContextVar
from typing import Optional, Dict, Any

# 定义上下文变量
user_id_context: ContextVar[Optional[str]] = ContextVar('user_id', default=None)
project_id_context: ContextVar[Optional[str]] = ContextVar('project_id', default=None)
request_id_context: ContextVar[Optional[str]] = ContextVar('request_id', default=None)
tenant_id_context: ContextVar[Optional[str]] = ContextVar('tenant_id', default=None)
user_info_context: ContextVar[Optional[Dict[str, Any]]] = ContextVar('user_info', default=None)

class RequestContext:
    """请求上下文管理器"""
    
    @staticmethod
    def set_user_id(user_id: str) -> None:
        """设置用户ID"""
        user_id_context.set(user_id)
    
    @staticmethod
    def get_user_id() -> Optional[str]:
        """获取用户ID"""
        return user_id_context.get()
    
    @staticmethod
    def set_project_id(project_id: str) -> None:
        """设置项目ID"""
        project_id_context.set(project_id)
    
    @staticmethod
    def get_project_id() -> Optional[str]:
        """获取项目ID"""
        return project_id_context.get()
    
    @staticmethod
    def set_request_id(request_id: str) -> None:
        """设置请求ID"""
        request_id_context.set(request_id)
    
    @staticmethod
    def get_request_id() -> Optional[str]:
        """获取请求ID"""
        return request_id_context.get()
    
    @staticmethod
    def set_tenant_id(tenant_id: str) -> None:
        """设置租户ID"""
        tenant_id_context.set(tenant_id)
    
    @staticmethod
    def get_tenant_id() -> Optional[str]:
        """获取租户ID"""
        return tenant_id_context.get()
    
    @staticmethod
    def set_user_info(user_info: Dict[str, Any]) -> None:
        """设置用户信息"""
        user_info_context.set(user_info)
    
    @staticmethod
    def get_user_info() -> Optional[Dict[str, Any]]:
        """获取用户信息"""
        return user_info_context.get()
    
    @staticmethod
    def clear() -> None:
        """清空上下文"""
        user_id_context.set(None)
        project_id_context.set(None)
        request_id_context.set(None)
        tenant_id_context.set(None)
        user_info_context.set(None)
    
    @staticmethod
    def get_context_info() -> dict:
        """获取所有上下文信息"""
        return {
            'user_id': RequestContext.get_user_id(),
            'project_id': RequestContext.get_project_id(),
            'request_id': RequestContext.get_request_id(),
            'tenant_id': RequestContext.get_tenant_id(),
            'user_info': RequestContext.get_user_info()
        }